import csv
import os
from SimTrader.tools import (debug_print, make_symbol_by_ins)
from SimTrader.objects import (MDFieldNames, Tick, Bar)
from SimTrader.constants import MDType
from SimTrader.simple_errors import MDTypeError


class CsvHelper(object):
    def __init__(self, file_path=''):
        self.__csv_reader = None
        self.__csv_writer = None
        self.__file_path = file_path

    def set_file_path(self, file_path):
        self.__file_path = file_path

    def load_row_as_dict(self, file_path="", headers=None, is_load_first_row=True):
        """
        load row as dict from csv
        :return: dict
        """
        if file_path:
            self.__file_path = file_path
        if not self.__csv_reader:
            self.__csv_reader = csv.reader(open(self.__file_path, 'r'))
        row_dict = {}
        is_first_row = True if is_load_first_row else False
        headers = list(headers) if headers else None
        for row in self.__csv_reader:
            if is_first_row and not headers:
                headers = row
                is_first_row = False
            else:
                for key, value in zip(headers, row):
                    row_dict[key] = value
                yield row_dict


class MDLoader(object):
    def __init__(self, base_path, fieldnames=None):
        self.__base_path = base_path
        self.__is_walk_base = True
        self.__md_helper = CsvHelper()
        if fieldnames:
            self.__fieldnames = fieldnames
        else:
            self.__fieldnames = MDFieldNames()

    def set_fieldnames(self, fieldnames):
        """
        设置加载文件的字段
        :param fieldnames: MDFieldName()
        :return: None
        """
        self.__fieldnames = fieldnames

    def set_file_path(self, file_path):
        """
        设置加载文件,若设置，只加载此文件的行情数据
        :param file_path: 文件目录
        :return: None
        """
        self.__md_helper.set_file_path(file_path)
        self.__is_walk_base = False

    def load_mkdata_from_csv(self, md_type, headers=None, file_path=None):
        """
        加载文件
        :param md_type: 加载的行情类型
        :param headers: 自定义行情文件字段
        :param file_path: 自定义行情文件
        :return: mkdata
        """
        if md_type == MDType.TICK:
            load_mkdata_func = self.__load_tick_from_csv
        elif md_type == MDType.BAR:
            load_mkdata_func = self.__load_bar_from_csv
        else:
            raise TypeError(MDTypeError)

        if not file_path and self.__is_walk_base:
            files = [os.path.join(self.__base_path, file) for file in os.listdir(self.__base_path)]
            for file in files:
                for mkdata in load_mkdata_func(file, headers):
                    yield mkdata
        else:
            for mkdata in load_mkdata_func(file_path, headers):
                yield mkdata

    def __load_tick_from_csv(self, file_path, headers):
        """
        加载tick行情
        :param file_path:行情文件
        :param headers: 行情文件字段
        :return: object<Tick>
        """
        for mkdata in self.__md_helper.load_row_as_dict(file_path=file_path, headers=headers):
            if not self.__fieldnames.trading_day:
                self.__fieldnames.trading_day = self.__fieldnames.date
            if not self.__fieldnames.action_day:
                self.__fieldnames.action_day = self.__fieldnames.date
            tick = Tick(
                trading_day=mkdata[self.__fieldnames.trading_day],
                action_day=mkdata[self.__fieldnames.action_day],
                time=mkdata[self.__fieldnames.time],
                open_price=float(mkdata[self.__fieldnames.open_price]),
                high_price=float(mkdata[self.__fieldnames.high_price]),
                low_price=float(mkdata[self.__fieldnames.low_price]),
                last_price=float(mkdata[self.__fieldnames.last_price]),
                volume=int(mkdata[self.__fieldnames.volume]),
                open_intl=float(mkdata[self.__fieldnames.open_intl]),
                turnover=float(mkdata.get(self.__fieldnames.turnover, 0.0)),
                avg_price=float(mkdata.get(self.__fieldnames.avg_Price, 0.0)),
                upper_limit=float(mkdata.get(self.__fieldnames.upper_limit, 0.0)),
                lower_limit=float(mkdata.get(self.__fieldnames.lower_limit, 0.0)),
                settelment=float(mkdata.get(self.__fieldnames.settlement, 0.0)),
                pre_close_price=float(mkdata.get(self.__fieldnames.pre_close_price, 0.0)),
                pre_settlement=float(mkdata.get(self.__fieldnames.pre_settlement, 0.0)),
                pre_open_intl=int(mkdata.get(self.__fieldnames.pre_open_intl, 0)),
                bid_price_1=float(mkdata.get(self.__fieldnames.bid_price_1, 0.0)),
                ask_price_1=float(mkdata.get(self.__fieldnames.ask_price_1, 0.0)),
                bid_volume_1=int(mkdata.get(self.__fieldnames.bid_volume_1, 0)),
                ask_volume_1=int(mkdata.get(self.__fieldnames.ask_volume_1, 0))
            )
            if not mkdata.get(self.__fieldnames.symbol):
                ins = mkdata[self.__fieldnames.instrument]
                tick.symbol = make_symbol_by_ins(ins)

            yield tick

    def __load_bar_from_csv(self, file_path, headers):
        """
        记载bar行情
        :param file_path: 行情文件
        :param headers: 行情文件的字段
        :return: object<Bar>
        """
        for mkdata in self.__md_helper.load_row_as_dict(file_path=file_path, headers=headers):
            bar = Bar(
                symbol=mkdata[self.__fieldnames.symbol],
                date=mkdata[self.__fieldnames.date],
                time=mkdata[self.__fieldnames.time],
                open_price=float(mkdata[self.__fieldnames.open_price]),
                high_price=float(mkdata[self.__fieldnames.high_price]),
                low_price=float(mkdata[self.__fieldnames.low_price]),
                close_price=float(mkdata[self.__fieldnames.close_price]),
                upper_limit=float(mkdata[self.__fieldnames.upper_limit]),
                lower_limit=float(mkdata[self.__fieldnames.lower_limit]),
                volume=int(mkdata[self.__fieldnames.volume]),
                open_intl=int(mkdata[self.__fieldnames.open_intl])
            )
            yield bar


if __name__ == "__main__":
    # csv_helper = CsvHelper()
    # csv_helper.set_file_path("C:\\Users\\zhzhrao_2020\\Desktop\\test_dict_loader.csv")
    # for row in csv_helper.load_row_as_dict():
    #     debug_print(row)

    md_loader = MDLoader("")
    md_loader.set_file_path("F:\\20200512\\ru2011.csv")
    fieldnames = MDFieldNames()
    md_loader.set_fieldnames(fieldnames)
    for tick in md_loader.load_mkdata_from_csv(MDType.TICK):
        debug_print(tick)